# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/8 10:18 
@Author  : ZhangShenao 
@File    : 1.使用自查询检索器实现动态元数据过滤.py
@Desc    : 使用自查询检索器实现动态元数据过滤

使用LangChain封装的SelfQueryRetriever自查询检索器,可以实现动态的元数据过滤。思路如下:
1. 通过Pydantic形式,定义元数据过滤器的结构
2. 利用LLM的Function Call功能,强制LLM进行函数调用,生成函数调用的结构化参数
3. 将生成的结构化参数,通过特定向量数据库的Translator翻译器,翻译成元数据过滤器
4. 使用VectorStoreRetriever进行检索,并插入元数据过滤器进行精确过滤

通过自查询的动态元数据过滤,可以有效压缩数据的检索范围,提升检索的准确率

自查询检索器对于面向特定领域的专用Agent效果相对较好,而对通用Agent来说效果较差
因为这些领域的文档一般相对来说比较规范,例如:财报、新闻、自媒体文章、教培等行业
这些行业的数据都能剥离出通用支持过滤与筛选的元数据/字段,使用自查询检索器能抽象出对应的检索字段信息
而对于通用的Agent,很难定义出通用的元数据用于筛选过滤。而如果元数据定义得过多,也会影响检索效果
"""
import os

import dotenv
from langchain.chains.query_constructor.schema import AttributeInfo
from langchain.retrievers import SelfQueryRetriever
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.documents import Document
from langchain_openai import ChatOpenAI
from langchain_pinecone import PineconeVectorStore

dotenv.load_dotenv()

# 构建文档列表,并保存到Pinecone向量数据库
docs = [
    Document(
        page_content="肖申克的救赎",
        metadata={"year": 1994, "rating": 9.7, "director": "弗兰克·德拉邦特"},
    ),
    Document(
        page_content="霸王别姬",
        metadata={"year": 1993, "rating": 9.6, "director": "陈凯歌"},
    ),
    Document(
        page_content="阿甘正传",
        metadata={"year": 1994, "rating": 9.5, "director": "罗伯特·泽米吉斯"},
    ),
    Document(
        page_content="泰坦尼克号",
        metadat={"year": 1997, "rating": 9.5, "director": "詹姆斯·卡梅隆"},
    ),
    Document(
        page_content="千与千寻",
        metadat={"year": 2001, "rating": 9.4, "director": "宫崎骏"},
    ),
    Document(
        page_content="星际穿越",
        metadat={"year": 2014, "rating": 9.4, "director": "克里斯托弗·诺兰"},
    ),
    Document(
        page_content="忠犬八公的故事",
        metadat={"year": 2009, "rating": 9.4, "director": "莱塞·霍尔斯道姆"},
    ),
    Document(
        page_content="三傻大闹宝莱坞",
        metadat={"year": 2009, "rating": 9.2, "director": "拉库马·希拉尼"},
    ),
    Document(
        page_content="疯狂动物城",
        metadat={"year": 2016, "rating": 9.2, "director": "拜伦·霍华德"},
    ),
    Document(
        page_content="无间道",
        metadat={"year": 2002, "rating": 9.3, "director": "刘伟强"},
    ),
]
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)
# vector_store.add_documents(docs)

# 定义子查询的元数据格式
metadata_field_info = [
    AttributeInfo(
        name="year",
        description="电影的发布年份",
        type="integer",
    ),
    AttributeInfo(
        name="rating",
        description="电影的评分",
        type="float",
    ),
    AttributeInfo(
        name="director",
        description="电影的导演",
        type="string",
    ),
]

# 创建SelfQueryRetriever子查询检索器
llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
retriever = SelfQueryRetriever.from_llm(
    llm=llm,  # 指定LLM大模型,用于执行Function Call
    vectorstore=vector_store,  # 指定向量数据库
    document_contents="电影的名称",  # 指定文档正文的描述
    metadata_field_info=metadata_field_info,  # 指定元数据字段格式
    enable_limit=True,  # 启用查询数量限制
)

# 执行检索
result = retriever.invoke("请帮我找出评分大于9.5的电影")
print(result)  # 肖申克的救赎+霸王别姬
print(len(result))  # 2

# 执行无元数据过滤的检索,会得到很多不相关的内容
# base_retriever = vector_store.as_retriever(
#     search_type="similarity",
#     search_kwargs={
#         "k": 2,
#     }
# )
# result = base_retriever.invoke("请帮我找出评分大于9.5的电影")
# print(result)
# print(len(result))
